"""
打印在模型中，某一层的特征
"""

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import pickle
import cv2


# 定义一个网络
class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=5, stride=1, padding=2),
            #  batch normalization
            nn.BatchNorm2d(16),
            # ReLU activation
            nn.ReLU(),
            # max pool
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.fc = nn.Linear(8 * 8 * 32, num_classes)

    # forward channel
    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out


# 模型训练
def train(model):
    # 随机数据
    x = torch.randn((32, 3, 32, 32))
    y = torch.randint(0, 10, (32,))

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    for it in range(3):
        out = model(x)
        loss = criterion(out, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # 保存模型
    model_info = {
        "iter_num": 3,
        "optimizer": optimizer,
        "model": model.state_dict()
    }
    torch.save(model_info, 'edu.ckpt')


# 获取模型参数
def get_params(model):
    params = {}
    for name, parameters in model.named_parameters():
        params[name] = parameters.detach().numpy()
    return params


def train_model_params(model):
    train(model)
    params = get_params(model)
    file_name = 'trainParams.pkl'

    with open(file_name, 'wb') as f:
        pickle.dump(params, f, pickle.HIGHEST_PROTOCOL)


def save_model_params(model):
    model_info = torch.load("edu.ckpt")
    model.load_state_dict(model_info["model"])
    params = get_params(model)
    file_name = 'saveParams.kpl'

    with open(file_name, 'wb') as f:
        pickle.dump(params, f, pickle.HIGHEST_PROTOCOL)


def test_params():
    model = ConvNet()
    train_model_params(model)
    save_model_params(model)

def show_params():
    train_file = 'trainParams.pkl'
    save_file = 'saveParams.kpl'
    with open(train_file, 'rb+') as f:
        train_params = pickle.load(f)

    with open(save_file, 'rb+') as f:
        save_params = pickle.load(f)

    print(type(train_params))
    print(train_params.keys())
    print(train_params['conv1.0.weight'])
    print(train_params['conv1.0.weight'].shape)
    print(train_params['conv1.0.weight'][0][0][0][0])
    print(save_params['conv1.0.weight'][0][0][0][0])


def test_image_in_layer():
    model = ConvNet()
    model_info = torch.load("model_info.ckpt")
    model.load_state_dict(model_info["model"])

    img = cv2.imread("cat.png")
    img = np.transpose(img, (2, 0, 1))
    img = np.expand_dims(img, 0)
    print(img.shape)
    img_tensor = torch.from_numpy(img).to(torch.float32)
    print(img_tensor)
    out = model.conv1[0](img_tensor)
    kernel = torch.ones((3, 16, 3, 3))

    out = torch.nn.functional.conv2d(out, kernel, bias=None, stride=1, padding=1, dilation=1, groups=1)

    out = out.to(torch.uint8).numpy()
    out = np.reshape(out, (3, 32, 32))
    print(out.shape)
    out = np.transpose(out, (1, 2, 0))

    print(out.shape)
    print(type(out))
    cv2.imshow("img", out)
    cv2.waitKey(0)

if __name__ == '__main__':
    # test_params()
    show_params()
    test_image_in_layer()